package com.affinity.impl.spi;

import java.util.BitSet;

import com.affinity.impl.AffinityMaskSpi;
import com.sun.jna.*;
import com.sun.jna.ptr.LongByReference;

/**
 * Currently support systems with <=64 cores -- since long used as bitmask for affinity set;
 * <p/>
 * <p/>
 * fixme: where are systems with more then 64 cores...
 *
 * @author cheremin
 * @since 28.04.12,  14:42
 */
public class PosixAffinityMaskSpi implements AffinityMaskSpi {

    /** @author BegemoT */
    private interface CLibrary extends Library {
        public static final CLibrary INSTANCE = ( CLibrary )
                Native.loadLibrary( "c", CLibrary.class );

        public int sched_setaffinity( final int pid,
                                      final int cpusetsize,
                                      final PointerType cpuset ) throws LastErrorException;

        public int sched_getaffinity( final int pid,
                                      final int cpusetsize,
                                      final PointerType cpuset ) throws LastErrorException;

        public int sched_getcpu() throws LastErrorException;
    }

    public PosixAffinityMaskSpi() {
    }

    @Override
    public void pingSpi() {
        final int cpusCount = Runtime.getRuntime().availableProcessors();

        if ( cpusCount > Long.SIZE ) {
            throw new UnsupportedOperationException( cpusCount + " > Long.SIZE -> not supported by current implementation" );
        }

        getCurrentThreadCPU();

        final BitSet affinityMask = getAffinityMask();
        setAffinityMask( affinityMask );
    }

    @Override
    public int getCurrentThreadCPU() {
        try {
            final int cpuNo = syslib().sched_getcpu();
            if ( cpuNo < 0 ) {
                throw new IllegalStateException( "sched_getcpu() return " + cpuNo );
            }
            return cpuNo;
        } catch ( LastErrorException e ) {
            throw new IllegalStateException( "sched_getcpu() errorNo=" + e.getErrorCode(), e );
        }
    }

    private static CLibrary syslib() {
        return CLibrary.INSTANCE;
    }

    @Override
    public BitSet getAffinityMask() {
        final LongByReference cpuset = new LongByReference( 0L );
        try {
            final int ret = syslib().sched_getaffinity(
                    0,//0 means "current thread id"
                    Long.SIZE / 8,//size in bytes
                    cpuset
            );
            if ( ret < 0 ) {
                throw new IllegalStateException( "sched_getaffinity((" + Long.SIZE / 8 + ") , &(" + cpuset + ") ) return " + ret );
            }
            return Utils.fromLong( cpuset.getValue() );
        } catch ( LastErrorException e ) {
            throw new IllegalStateException( "sched_getaffinity((" + Long.SIZE / 8 + ") , &(" + cpuset + ") ) errorNo=" + e.getErrorCode(), e );
        }
    }

    @Override
    public void setAffinityMask( final BitSet affinityMask ) {
        if ( affinityMask == null || affinityMask.isEmpty() ) {
            throw new IllegalArgumentException( "affinityMask can't be null nor empty" );
        }
        try {
            final long longMask = Utils.toLong( affinityMask );
            final int ret = syslib().sched_setaffinity(
                    0,              //0 means "current thread id"
                    Long.SIZE / 8,  //== size in bytes
                    new LongByReference( longMask )
            );
            if ( ret < 0 ) {
                throw new IllegalStateException( "sched_setaffinity((" + Long.SIZE / 8 + ") , &(" + affinityMask + ") ) return " + ret );
            }
        } catch ( LastErrorException e ) {
            throw new IllegalStateException( "sched_setaffinity((" + Long.SIZE / 8 + ") , &(" + affinityMask + ") ) errorNo=" + e.getErrorCode(), e );
        }
    }
}
